#include "Command.h"
#include "CConsole.h"

void ClearScreen(const std::string& strInput, const std::string& strCommand, const std::vector<std::string>& argv, void* lpUserPtr, BOOL& bExit)
{
    CConsole::GetInstance()->ClearScreen();
}

void Initialize(const std::string& strInput, const std::string& strCommand, const std::vector<std::string>& argv, void* lpUserPtr, BOOL& bExit)
{
    PHANDLE lphFile = (PHANDLE)lpUserPtr;
    *lphFile = CreateFile(USTRING_WIN32_NAME, GENERIC_ALL, 0, nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (INVALID_HANDLE_VALUE == *lphFile)
        CConsole::OutputStringPrint("Initialize Error\r\n");
    else
        CConsole::OutputStringPrint("Initialize Success\r\n");
}

void Quit(const std::string& strInput, const std::string& strCommand, const std::vector<std::string>& argv, void* lpUserPtr, BOOL& bExit)
{
    if (lpUserPtr)
    {
        PHANDLE lphFile = (PHANDLE)lpUserPtr;
        if (*lphFile && INVALID_HANDLE_VALUE == *lphFile)
            CloseHandle(*lphFile);
    }

    bExit = TRUE;
}

void InjectTo(const std::string& strInput, const std::string& strCommand, const std::vector<std::string>& argv, void* lpUserPtr, BOOL& bExit)
{
    HANDLE hFile = *((PHANDLE)lpUserPtr);
    ULONG dwProcessId;
    BOOL bSuccess = FALSE;
    std::stringstream ss;
    IOTYPE_PROCESS_INJECT_INFO InjectInfo = { 0 };
    WCHAR wstrFileName[MAX_PATH] = { 0 };
    DWORD dwLength;
    PWCH lpWch;
    std::wstring str;
    HANDLE hProcess = nullptr;
    BOOL bWow64Process;
    BOOL bRet;
    DWORD dwRet;

    do
    {
        if (argv.size() != 1) break;

        ss << argv[0];
        if (argv[0].find("0x") != std::string::npos || argv[0].find("0X") != std::string::npos)
            ss >> std::hex >> dwProcessId;
        else
            ss >> dwProcessId;

        hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwProcessId);
        if (!hProcess) break;

        if (!IsWow64Process(hProcess, &bWow64Process)) break;

        dwLength = GetModuleFileName(nullptr, wstrFileName, MAX_PATH);
        if (!dwLength) break;

        lpWch = wstrFileName + dwLength;
        while (lpWch != wstrFileName && *lpWch != L'\\') lpWch--;
        if (lpWch == wstrFileName) break;
        lpWch++;
        *lpWch = L'\0';
        str = wstrFileName;

        if (bWow64Process)
            str += L"Test\\Dll32.dll";
        else
            str += L"Test\\Dll64.dll";

        InjectInfo.ProcessId = dwProcessId;
        InjectInfo.Bytes = str.size() * sizeof(WCHAR);
        InjectInfo.lpModuleName = (PVOID)(str.c_str());
        bRet = DeviceIoControl(
            hFile,
            IOCTL_INJECT_TO_PROCESS,
            &InjectInfo,
            sizeof(IOTYPE_PROCESS_INJECT_INFO),
            nullptr,
            0,
            &dwRet,
            nullptr
        );
        if (!bRet) break;

        bSuccess = TRUE;

    } while (FALSE);

    if (hProcess)
        CloseHandle(hProcess);

    if (bSuccess)
        CConsole::OutputStringPrint(std::string("Inject Success\r\n"));
    else
        CConsole::OutputStringPrint(std::string("Inject Error\r\n"));
}
